import subprocess
import re
import cv2
import requests
import base64
import json


# TPU供应商ID和产品ID
tpu_dict = {
    "kneron":   "3231:0100",
    "songke":   "0525:a4a0",
    "amc":      "0329:4c47",
    "K720":     "3231:0200"
}


# 检测TPU类型
def check_tpu_type(device_dict):

    tpu_type_list = []
    try:
        # Execute lsusb command
        output = subprocess.check_output(["lsusb"], text=True)

        # Parse each line of the lsusb output
        for line in output.splitlines():
            match = re.search(r"ID (\w+:\w+)", line)
            if match:
                detected_id = match.group(1)
                # Check if the detected ID is in the given dictionary values
                for device_name, vendor_product_id in device_dict.items():
                    if detected_id == vendor_product_id:
                        tpu_type_list.append(device_name)
                        break  # Avoid duplicate entries
    except subprocess.CalledProcessError as e:
        print(f"Error running lsusb: {e}")
    except Exception as e:
        print(f"An error occurred: {e}")
    
    return tpu_type_list


class TPU_():
    def __init__(self):

        self.tpu_init_flag = 0
        

    # 获取图片
    def get_img(self):
        from cap import VideoCapture
        try:
            cap = VideoCapture()
        except Exception as e:
            print(f"初始化摄像头异常, error: {e}")
            exit(0)
        image = cap.read()
        cap.release()
        if image is not None:
            return image

    def predict(self, image, mode=True, model_name="", iou=0.45, conf=0.1, w=640, h=640, url = '192.168.137.1'):
        # 如果是联调
        if mode == True:
            _, img_encoded = cv2.imencode('.png', image)
            # 将编码后的图像转换为 Base64 字符串
            img_base64 = base64.b64encode(img_encoded).decode('utf-8')
            # 构造请求数据
            data = {
            "data": img_base64,
            "modelName": model_name,
            "confThres": conf,
            "iouThres": iou,
            "width": w,
            "height": h
            }
            # 发送 POST 请求
            url = "http://" + url + ":5500/predict"
            response = requests.post(url, json=data)
            # 判断响应是否成功
            if response.status_code == 200:
                # 获取响应内容
                result = response.json()
                if result is not None:
                    # print(f"result:{result}")
                    box =[]
                    for l1 in result['data']:
                        xy1 = l1['left_up']
                        xy2 = l1['right_down']
                        best_result = max(l1["data"], key=lambda x: x['confidence'])
                        label = best_result['tag_name']
                        con = best_result['confidence']
                        w = xy2[0] - xy1[0]
                        h = xy2[1] - xy1[1]
                        pixelX = xy1[0] + w/2
                        pixelY = xy1[1] + h/2
                        box.append([int(pixelX), int(pixelY), int(w), int(h), con, label])
                    if box is not None:
                        # print(f"box:{box}")
                        return box
                else:
                    print("联调结果为空")
            else:
                print(f"Request failed with status code {response.status_code}")
                return False
        else:
            results_box = []
            if self.tpu_init_flag == 0:
                # TPU的类型检测，采用检测其供应商ID和产品ID的方式
                tpu_type_list = check_tpu_type(tpu_dict)
                if len(tpu_type_list) == 0:
                    print("未检测到加速棒")
                    exit(0)
                elif len(tpu_type_list) == 1:
                    self.tpu_type = tpu_type_list[0]
                    print('tpu_type : ', self.tpu_type)
                from TPU import TPUFactory
                print("成功导入TPU库")

                if model_name in ["static_grab", "refuse_classify", "calculate"]:
                    self.tpu = TPUFactory(algorithm = "yolov5")
                elif model_name in ["tictactoe"]:
                    pass
                else:
                    print("请填写正确的模型名称")
                    exit(0)

                if self.tpu.link() == False:
                    print(f"初始化TPU异常")
                    exit(0)

                # 获取模型
                model_path = "/opt/SONGKE_BOT/tpu/model"
                if self.tpu_type == "kneron":
                    self.tpu.model_init(model_path)
                elif self.tpu_type == "songke":
                    res = self.tpu.get_model(model_name, model_path)
                    if res:
                        print("模型下载成功,路径为:", model_path)
                    else:
                        print(f"模型下载错误")
                        exit(0)
                    self.tpu.model_init(model_path, host=True)
                self.tpu_init_flag = 1

            box = self.tpu.predict(image, [w,h], conf, iou)

            if box is not None:
                for o in box:
                    x = int(o[0] / 640 * w)
                    y = int(o[1] / 480 * h)
                    w_ = int(o[2] / 640 * w)
                    h_ = int(o[3] / 480 * h)

                    # 将处理后的坐标、宽高和标签存储到列表中
                    results_box.append([int(x), int(y), int(w_), int(h_), o[4], o[5]])
                return results_box

    def __del__(self):
        # 确保资源在对象销毁时释放
        if hasattr(self, "tpu") and self.tpu is not None:
            try:
                self.tpu.release()
            except Exception as e:
                print(f"释放TPU资源时出错: {e}")


Control_TPU = TPU_()

# image = Control_TPU.get_img()

# box = Control_TPU.predict(image, mode=False, model_name="static_grab", iou=0.45, conf=0.65, w=416, h=416)

# print("获取完成")
